/******************************************************************************
 * Copyright 2017 The Apollo Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/

/**
 * @file
 */

#ifndef MODULES_DREAMVIEW_BACKEND_WEBSOCKET_H_
#define MODULES_DREAMVIEW_BACKEND_WEBSOCKET_H_

#include <memory>
#include <mutex>
#include <set>
#include <string>
#include <thread>
#include <unordered_map>

#include "CivetServer.h"
#include "third_party/json/json.hpp"

/**
 * @namespace apollo::dreamview
 * @brief apollo::dreamview
 */
namespace apollo {
namespace dreamview {

/**
 * @class WebSocketHandler
 *
 * @brief The WebSocketHandler, built on top of CivetWebSocketHandler,
 * is a websocket handler that handles different types of websocket related
 * events.
 * the server and a client endpoint. The SendData() method is used to push data
 * to all the connected clients.
 */
class WebSocketHandler : public CivetWebSocketHandler {
 public:
  using Json = nlohmann::json;
  using Connection = struct mg_connection;
  using MessageHandler = std::function<void(const Json &, Connection *)>;

  /**
   * @brief Callback method for when the client intends to establish a websocket
   * connection, before websocket handshake.
   *
   * @param server the calling server
   * @param conn the connection information
   * @returns true to keep socket open, false to close it
   */
  bool handleConnection(CivetServer *server, const Connection *conn) override {
    return true;
  }

  /**
   * @brief Callback method for when websocket handshake is successfully
   * completed, and connection is ready for data exchange.
   *
   * @param server the calling server
   * @param conn the connection information
   */
  void handleReadyState(CivetServer *server, Connection *conn) override;

  /**
   * @brief Callback method for when a data frame has been received from the
   * client.
   *
   * @param server the calling server
   * @param conn the connection information
   * @param bits first byte of the websocket frame, see websocket RFC at
   *             http://tools.ietf.org/html/rfc6455, section 5.2
   * @param data payload, with mask (if any) already applied.
   * @param data_len length of data
   * @returns true to keep socket open, false to close it
   */
  bool handleData(CivetServer *server, Connection *conn, int bits, char *data,
                  size_t data_len) override;

  /**
   * @brief Callback method for when the connection is closed.
   *
   * @param server the calling server
   * @param conn the connection information
   */
  void handleClose(CivetServer *server, const Connection *conn) override;

  /**
   * @brief Sends the provided data to all the connected clients.
   * @param data The message string to be sent.
   */
  bool SendData(const std::string &data);
  bool SendData(const std::string &data, Connection *conn);

  /**
   * @brief Add a new message handler for a message type.
   * @param type The name/key to identify the message type.
   * @param handler The function to handle the received message.
   */
  void RegisterMessageHandler(std::string type, MessageHandler handler) {
    message_handlers_[type] = handler;
  }

 private:
  // Message handlers keyed by message type.
  std::unordered_map<std::string, MessageHandler> message_handlers_;

  // The mutex guarding the connection set. We are not using read
  // write lock, as the server is not expected to get many clients
  // (connections).
  mutable std::mutex mutex_;

  // The pool of all maintained connections.
  std::set<Connection *> connections_;
};

}  // namespace dreamview
}  // namespace apollo

#endif /* MODULES_DREAMVIEW_BACKEND_WEBSOCKET_H_ */
